#include "vio/tools/homography.h"

namespace vio {
namespace vision_utils {

Homography::Homography(
    const std::vector<Eigen::Vector2d,
                      Eigen::aligned_allocator<Eigen::Vector2d>>& fts1,
    const std::vector<Eigen::Vector2d,
                      Eigen::aligned_allocator<Eigen::Vector2d>>& fts2,
    double error_multiplier2,
    double thresh_in_px)
    : thresh_(thresh_in_px),
      error_multiplier2_(error_multiplier2),
      fts_c1_(fts1),
      fts_c2_(fts2) {}

// ref: https://docs.opencv.org/3.4/d9/dab/tutorial_homography.html
void Homography::CalcFromPlaneParams(const Eigen::Vector3d& n_c1,
                                     const Eigen::Vector3d& xyz_c1) {
  // normal distance from plane to KF
  double d = n_c1.dot(xyz_c1);
  H_c2_from_c1_ = T_c2_from_c1_.rotationMatrix() +
                  (T_c2_from_c1_.translation() * n_c1.transpose()) / d;
}

void Homography::CalcFromMatches() {
  std::vector<cv::Point2f> src_pts(fts_c1_.size()), dts_pts(fts_c1_.size());
  for (size_t i = 0; i < fts_c1_.size(); ++i) {
    src_pts[i] = cv::Point2f(fts_c1_[i][0], fts_c1_[i][1]);
    dts_pts[i] = cv::Point2f(fts_c2_[i][0], fts_c2_[i][1]);
  }

  // TODO: replace this function to remove dependency from opencv!
  cv::Mat cvH = cv::findHomography(
      src_pts, dts_pts, cv::RANSAC, 2.0 / error_multiplier2_);
  cv::cv2eigen(cvH, H_c2_from_c1_);
}

size_t Homography::ComputeMatchesIniers() {
  inliers_.clear();
  inliers_.resize(fts_c1_.size());
  size_t n_inliers = 0;
  for (size_t i = 0; i < fts_c1_.size(); i++) {
    // c2_pixel = K * H * K_inv * c1_pixel
    Eigen::Vector2d projected = math_utils::project2d(
        H_c2_from_c1_ * math_utils::unproject2d(fts_c1_[i]));
    Eigen::Vector2d e = fts_c2_[i] - projected;
    double e_px = error_multiplier2_ * e.norm();
    inliers_[i] = (e_px < thresh_);
    n_inliers += inliers_[i];
  }
  return n_inliers;
}

bool Homography::ComputeSE3FromMatches() {
  CalcFromMatches();
  bool res = Decompose();
  if (!res) {
    return false;
  }
  ComputeMatchesIniers();
  FindBestDecomposition();
  T_c2_from_c1_ = decompositions_.front().T;
  return true;
}

// ref: Motion and Structure from Motion in a Piecewise Planar Environment
// https://hal.inria.fr/file/index/docid/75698/filename/RR-0856.pdf
bool Homography::Decompose() {
  decompositions_.clear();
  Eigen::JacobiSVD<Eigen::MatrixXd> svd(
      H_c2_from_c1_, Eigen::ComputeThinU | Eigen::ComputeThinV);
  Eigen::Vector3d singular_values = svd.singularValues();

  double d1 =
      std::fabs(singular_values[0]);  // The paper suggests the square of these
                                      // (e.g. the evalues of AAT)
  double d2 = std::fabs(singular_values[1]);  // should be used, but this is
                                              // wrong. c.f Faugeras' book
  double d3 = std::fabs(singular_values[2]);

  Eigen::Matrix3d U = svd.matrixU();
  Eigen::Matrix3d V = svd.matrixV();

  double s = U.determinant() * V.determinant();

  double dPrime_PM = d2;

  int nCase;
  if (d1 != d2 && d2 != d3) {
    nCase = 1;
  } else if (d1 == d2 && d2 == d3) {
    nCase = 3;
  } else {
    nCase = 2;
  }

  if (nCase != 1) {
    LOG(INFO) << LOG_RED_BEGIN
              << "FATAL Homography Initialization: This motion case is not "
                 "implemented or is degenerate. Try again."
              << LOG_COLOR_END;
    return false;
  }

  double x1_PM;
  double x2;
  double x3_PM;

  // all below deals with the case = 1 case.
  // case 1 implies(d1 != d3)
  {  // Eq. 12
    x1_PM = std::sqrt((d1 * d1 - d2 * d2) / (d1 * d1 - d3 * d3));
    x2 = 0.0;
    x3_PM = std::sqrt((d2 * d2 - d3 * d3) / (d1 * d1 - d3 * d3));
  }
  double e1[4] = {1.0, -1.0, 1.0, -1.0};
  double e3[4] = {1.0, 1.0, -1.0, -1.0};

  Eigen::Vector3d np;
  HomographyDecomposition decomp;

  // case 1, d' > 0:
  decomp.d = s * dPrime_PM;
  for (size_t signs = 0; signs < 4; signs++) {
    // Eq 13
    decomp.R = Eigen::Matrix3d::Identity();
    double dSinTheta = (d1 - d3) * x1_PM * x3_PM * e1[signs] * e3[signs] / d2;
    double dCosTheta = (d1 * x3_PM * x3_PM + d3 * x1_PM * x1_PM) / d2;
    decomp.R(0, 0) = dCosTheta;
    decomp.R(0, 2) = -dSinTheta;
    decomp.R(2, 0) = dSinTheta;
    decomp.R(2, 2) = dCosTheta;

    // Eq 14
    decomp.t[0] = (d1 - d3) * x1_PM * e1[signs];
    decomp.t[1] = 0.0;
    decomp.t[2] = (d1 - d3) * -x3_PM * e3[signs];

    np[0] = x1_PM * e1[signs];
    np[1] = x2;
    np[2] = x3_PM * e3[signs];
    decomp.n = V * np;

    decompositions_.push_back(decomp);
  }

  // Case 1, d' < 0:
  decomp.d = s * -dPrime_PM;
  for (size_t signs = 0; signs < 4; signs++) {
    // Eq 15
    decomp.R = -1 * Eigen::Matrix3d::Identity();
    double dSinPhi = (d1 + d3) * x1_PM * x3_PM * e1[signs] * e3[signs] / d2;
    double dCosPhi = (d3 * x1_PM * x1_PM - d1 * x3_PM * x3_PM) / d2;
    decomp.R(0, 0) = dCosPhi;
    decomp.R(0, 2) = dSinPhi;
    decomp.R(2, 0) = dSinPhi;
    decomp.R(2, 2) = -dCosPhi;

    // Eq 16
    decomp.t[0] = (d1 + d3) * x1_PM * e1[signs];
    decomp.t[1] = 0.0;
    decomp.t[2] = (d1 + d3) * x3_PM * e3[signs];

    np[0] = x1_PM * e1[signs];
    np[1] = x2;
    np[2] = x3_PM * e3[signs];
    decomp.n = V * np;

    decompositions_.push_back(decomp);
  }

  // save rotation and translation of the decomposition
  for (unsigned int i = 0; i < decompositions_.size(); i++) {
    Eigen::Matrix3d R = s * U * decompositions_[i].R * V.transpose();
    Eigen::Vector3d t = U * decompositions_[i].t;
    decompositions_[i].T = Sophus::SE3d(R, t);
  }
  return true;
}

bool operator<(const HomographyDecomposition lhs,
               const HomographyDecomposition rhs) {
  return lhs.score < rhs.score;
}

void Homography::FindBestDecomposition() {
  assert(decompositions_.size() == 8);
  for (size_t i = 0; i < decompositions_.size(); i++) {
    HomographyDecomposition& decom = decompositions_[i];
    size_t nPositive = 0;
    for (size_t m = 0; m < fts_c1_.size(); m++) {
      if (!inliers_[m])
        continue;
      const Eigen::Vector2d& v2 = fts_c1_[m];
      double dVisibilityTest =
          (H_c2_from_c1_(2, 0) * v2[0] + H_c2_from_c1_(2, 1) * v2[1] +
           H_c2_from_c1_(2, 2)) /
          decom.d;
      if (dVisibilityTest > 0.0)
        nPositive++;
    }
    decom.score = -nPositive;
  }

  std::sort(decompositions_.begin(), decompositions_.end());
  decompositions_.resize(4);

  for (size_t i = 0; i < decompositions_.size(); i++) {
    HomographyDecomposition& decom = decompositions_[i];
    int nPositive = 0;
    for (size_t m = 0; m < fts_c1_.size(); m++) {
      if (!inliers_[m])
        continue;
      Eigen::Vector3d v3 = math_utils::unproject2d(fts_c1_[m]);
      double dVisibilityTest = v3.dot(decom.n) / decom.d;
      if (dVisibilityTest > 0.0)
        nPositive++;
    };
    decom.score = -nPositive;
  }

  std::sort(decompositions_.begin(), decompositions_.end());
  decompositions_.resize(2);

  // According to Faugeras and Lustman, ambiguity exists if the two scores are
  // equal but in practive, better to look at the ratio!
  double dRatio =
      (double)decompositions_[1].score / (double)decompositions_[0].score;

  if (dRatio < 0.9)  // no ambiguity!
    decompositions_.erase(decompositions_.begin() + 1);
  else  // two-way ambiguity. Resolve by sampsonus score of all points.
  {
    double dErrorSquaredLimit = thresh_ * thresh_ * 4;
    double adSampsonusScores[2];
    for (size_t i = 0; i < 2; i++) {
      Sophus::SE3d T = decompositions_[i].T;
      Eigen::Matrix3d Essential =
          T.rotationMatrix() * math_utils::sqew(T.translation());
      double dSumError = 0;
      for (size_t m = 0; m < fts_c1_.size(); m++) {
        double d =
            math_utils::SampsonusError(fts_c1_[m], Essential, fts_c2_[m]);
        if (d > dErrorSquaredLimit)
          d = dErrorSquaredLimit;
        dSumError += d;
      }
      adSampsonusScores[i] = dSumError;
    }

    if (adSampsonusScores[0] <= adSampsonusScores[1])
      decompositions_.erase(decompositions_.begin() + 1);
    else
      decompositions_.erase(decompositions_.begin());
  }
}
}  // namespace vision_utils
}  // namespace vio